#!/bin/bash
# ==============================================================================
# GSICoreAnalysis.jl - Installation and Setup Script
# ==============================================================================
# This script handles complete installation and setup of the GSI package

set -e

# Color codes
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
MAGENTA='\033[0;35m'
NC='\033[0m'

# Configuration
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
JULIA_MIN_VERSION="1.8"
REQUIRED_MEMORY_GB=4

print_banner() {
    echo -e "${MAGENTA}================================================${NC}"
    echo -e "${MAGENTA}  GSICoreAnalysis.jl - Setup & Installation${NC}"
    echo -e "${MAGENTA}================================================${NC}"
    echo "🌍 Atmospheric Data Assimilation Package Setup"
    echo ""
    echo "📁 Installation Directory: $SCRIPT_DIR"
    echo "📋 Minimum Julia Version: $JULIA_MIN_VERSION"
    echo "💾 Required Memory: ${REQUIRED_MEMORY_GB}GB"
    echo ""
}

print_usage() {
    echo "Usage: $0 [OPTIONS] [SETUP_TYPE]"
    echo ""
    echo "SETUP TYPES:"
    echo "  full           - Complete installation with all features"
    echo "  minimal        - Minimal installation (core functionality only)"
    echo "  development    - Development setup with all tools"
    echo "  visualization  - Setup with visualization capabilities"
    echo "  hpc            - High-performance computing configuration"
    echo ""
    echo "OPTIONS:"
    echo "  -h, --help     - Show this help message"
    echo "  -y, --yes      - Automatic yes to all prompts"
    echo "  -f, --force    - Force reinstallation"
    echo "  --no-test      - Skip post-installation tests"
    echo "  --no-precompile- Skip package precompilation"
    echo "  --threads N    - Set default Julia threads (default: auto)"
    echo "  --check-only   - Only check system requirements"
    echo ""
    echo "EXAMPLES:"
    echo "  $0                  # Interactive full installation"
    echo "  $0 -y minimal       # Automatic minimal setup"
    echo "  $0 --check-only     # Check system requirements only"
    echo "  $0 development      # Development setup"
    echo ""
}

check_system_requirements() {
    echo -e "${BLUE}🔍 Checking system requirements...${NC}"

    local errors=0

    # Check operating system
    case "$(uname)" in
        "Linux")
            echo -e "${GREEN}✓ Operating System: Linux${NC}"
            ;;
        "Darwin")
            echo -e "${GREEN}✓ Operating System: macOS${NC}"
            ;;
        *)
            echo -e "${RED}❌ Unsupported OS: $(uname)${NC}"
            ((errors++))
            ;;
    esac

    # Check memory
    local memory_gb
    if command -v free >/dev/null 2>&1; then
        memory_gb=$(free -g | awk 'NR==2{print $2}')
    elif [[ "$(uname)" == "Darwin" ]]; then
        memory_gb=$(($(sysctl -n hw.memsize) / 1024 / 1024 / 1024))
    else
        memory_gb=0
    fi

    if [[ $memory_gb -ge $REQUIRED_MEMORY_GB ]]; then
        echo -e "${GREEN}✓ Memory: ${memory_gb}GB (>= ${REQUIRED_MEMORY_GB}GB required)${NC}"
    else
        echo -e "${YELLOW}⚠️ Memory: ${memory_gb}GB (${REQUIRED_MEMORY_GB}GB recommended)${NC}"
    fi

    # Check disk space
    local disk_gb=$(df -BG "$SCRIPT_DIR" | awk 'NR==2 {print $4}' | sed 's/G//')
    if [[ $disk_gb -ge 2 ]]; then
        echo -e "${GREEN}✓ Disk Space: ${disk_gb}GB available${NC}"
    else
        echo -e "${RED}❌ Disk Space: ${disk_gb}GB (2GB required)${NC}"
        ((errors++))
    fi

    return $errors
}

check_julia_installation() {
    echo -e "${BLUE}🔍 Checking Julia installation...${NC}"

    if ! command -v julia >/dev/null 2>&1; then
        echo -e "${RED}❌ Julia not found in PATH${NC}"
        echo ""
        echo "Please install Julia:"
        echo "  • Download from: https://julialang.org/downloads/"
        echo "  • Or use package manager:"
        echo "    - Ubuntu: sudo snap install julia --classic"
        echo "    - macOS: brew install julia"
        echo ""
        return 1
    fi

    local julia_version=$(julia --version | grep -oP 'version \K[0-9]+\.[0-9]+')
    local required_version="$JULIA_MIN_VERSION"

    if [[ "$(printf '%s\n' "$required_version" "$julia_version" | sort -V | head -n1)" == "$required_version" ]]; then
        echo -e "${GREEN}✓ Julia version: $julia_version (>= $required_version)${NC}"
    else
        echo -e "${RED}❌ Julia version: $julia_version (>= $required_version required)${NC}"
        return 1
    fi

    # Check Julia package manager
    if julia -e "using Pkg" >/dev/null 2>&1; then
        echo -e "${GREEN}✓ Julia package manager available${NC}"
    else
        echo -e "${RED}❌ Julia package manager not working${NC}"
        return 1
    fi

    return 0
}

setup_julia_environment() {
    local setup_type="$1"

    echo -e "${BLUE}📦 Setting up Julia environment ($setup_type)...${NC}"

    cd "$SCRIPT_DIR"

    # Check Project.toml exists
    if [[ ! -f "Project.toml" ]]; then
        echo -e "${RED}❌ Project.toml not found in $SCRIPT_DIR${NC}"
        return 1
    fi

    echo "📋 Installing package dependencies..."

    # Install core dependencies
    julia --project=. -e "
    using Pkg
    println(\"📦 Installing core dependencies...\")
    Pkg.instantiate()
    println(\"✓ Core dependencies installed\")
    " || {
        echo -e "${RED}❌ Failed to install core dependencies${NC}"
        return 1
    }

    # Setup type-specific dependencies
    case "$setup_type" in
        "full"|"development"|"visualization")
            echo "📊 Installing visualization dependencies..."
            julia --project=. -e "
            using Pkg
            viz_packages = [\"Plots\", \"StatsPlots\", \"PlotlyJS\", \"StatsBase\", \"Distributions\"]
            Pkg.add(viz_packages)
            println(\"✓ Visualization dependencies installed\")
            " || {
                echo -e "${YELLOW}⚠️ Some visualization packages may not be available${NC}"
            }
            ;;
        "hpc")
            echo "🚀 Installing HPC dependencies..."
            julia --project=. -e "
            using Pkg
            hpc_packages = [\"MPI\", \"CUDA\"]
            try
                Pkg.add(hpc_packages)
                println(\"✓ HPC dependencies installed\")
            catch e
                println(\"⚠️ Some HPC packages not available: \", e)
            end
            "
            ;;
    esac

    # Precompile packages
    if [[ "$SKIP_PRECOMPILE" != "true" ]]; then
        echo "🔨 Precompiling packages..."
        julia --project=. -e "
        using Pkg
        println(\"🔨 Precompiling packages...\")
        Pkg.precompile()
        println(\"✓ Precompilation completed\")
        " || {
            echo -e "${YELLOW}⚠️ Precompilation had issues but may still work${NC}"
        }
    fi

    echo -e "${GREEN}✓ Julia environment setup completed${NC}"
}

create_run_scripts_executable() {
    echo -e "${BLUE}🔧 Making run scripts executable...${NC}"

    local scripts=(
        "run_gsi_analysis.sh"
        "run_gsi_visualization.sh"
        "run_gsi_tests.sh"
        "setup_gsi.sh"
    )

    for script in "${scripts[@]}"; do
        if [[ -f "$SCRIPT_DIR/$script" ]]; then
            chmod +x "$SCRIPT_DIR/$script"
            echo -e "${GREEN}✓ $script${NC}"
        fi
    done
}

setup_environment_file() {
    echo -e "${BLUE}📄 Creating environment configuration...${NC}"

    local env_file="$SCRIPT_DIR/.gsi_env"

    local overlay_dir="$SCRIPT_DIR/deps/eccodes_overrides"

    cat > "$env_file" << EOF
# GSICoreAnalysis.jl Environment Configuration
# Generated: $(date)

export GSI_ROOT="$SCRIPT_DIR"
export GSI_LOG_DIR="$SCRIPT_DIR/logs"
export GSI_RESULTS_DIR="$SCRIPT_DIR/results"
export GSI_VIS_DIR="$SCRIPT_DIR/drp4dvar_visualizations"

# Julia configuration
export JULIA_PROJECT="$SCRIPT_DIR"
export JULIA_NUM_THREADS=${JULIA_THREADS:-auto}

# GSI specific settings
export GSI_PLOT_BACKEND="plotlyjs"
export GSI_PLOT_DPI="300"
export GSI_MEMORY_LIMIT="4GB"

# Convenience aliases
alias gsi-run="$SCRIPT_DIR/run_gsi_analysis.sh"
alias gsi-viz="$SCRIPT_DIR/run_gsi_visualization.sh"
alias gsi-test="$SCRIPT_DIR/run_gsi_tests.sh"
alias gsi-cd="cd $SCRIPT_DIR"

# Usage information
gsi-help() {
    echo "GSICoreAnalysis.jl - Available Commands:"
    echo "  gsi-run [type]    - Run GSI analysis"
    echo "  gsi-viz [type]    - Run with visualization"
    echo "  gsi-test [type]   - Run tests"
    echo "  gsi-cd            - Change to GSI directory"
    echo ""
    echo "To load this environment: source $env_file"
}

# ecCodes overlay configuration
export ECCODES_OVERLAY_DIR="$overlay_dir"

if [ -d "$ECCODES_OVERLAY_DIR" ]; then
    _gsi_defs="$ECCODES_OVERLAY_DIR"
    if [ -n "${ECCODES_HOME:-}" ] && [ -d "${ECCODES_HOME}/share/eccodes/definitions" ]; then
        _gsi_defs="$_gsi_defs:${ECCODES_HOME}/share/eccodes/definitions"
    fi
    if [ -n "${ECCODES_DEFINITION_PATH:-}" ]; then
        _gsi_defs="$_gsi_defs:${ECCODES_DEFINITION_PATH}"
    fi
    export ECCODES_DEFINITION_PATH="$_gsi_defs"
    unset _gsi_defs

    _gsi_samples=""
    if [ -n "${ECCODES_HOME:-}" ] && [ -d "${ECCODES_HOME}/share/eccodes/samples" ]; then
        _gsi_samples="${ECCODES_HOME}/share/eccodes/samples"
    fi
    if [ -n "${ECCODES_SAMPLES_PATH:-}" ]; then
        if [ -n "$_gsi_samples" ]; then
            _gsi_samples="$_gsi_samples:${ECCODES_SAMPLES_PATH}"
        else
            _gsi_samples="${ECCODES_SAMPLES_PATH}"
        fi
    fi
    if [ -n "$_gsi_samples" ]; then
        export ECCODES_SAMPLES_PATH="$_gsi_samples"
    fi
    unset _gsi_samples
fi
EOF

    echo -e "${GREEN}✓ Environment file created: $env_file${NC}"
    echo "💡 To use: source $env_file"
}

run_post_install_tests() {
    if [[ "$SKIP_TESTS" == "true" ]]; then
        echo -e "${YELLOW}⚠️ Skipping post-installation tests${NC}"
        return 0
    fi

    echo -e "${BLUE}🧪 Running post-installation tests...${NC}"

    # Test basic functionality
    echo "Testing basic GSI functionality..."

    local test_script=$(mktemp --suffix=.jl)
    cat > "$test_script" << 'EOF'
# Basic GSI test
try
    using GSICoreAnalysis
    println("✓ GSICoreAnalysis module loaded successfully")

    # Test basic configuration
    config = AnalysisConfig(grid_size=(10, 10, 5), ensemble_size=2)
    println("✓ AnalysisConfig creation successful")

    println("✓ Basic GSI functionality test passed")
catch e
    println("❌ Basic GSI functionality test failed: ", e)
    exit(1)
end
EOF

    if julia --project=. "$test_script"; then
        echo -e "${GREEN}✓ Basic functionality test passed${NC}"
    else
        echo -e "${RED}❌ Basic functionality test failed${NC}"
        rm -f "$test_script"
        return 1
    fi

    rm -f "$test_script"

    # Test quick demo if available
    if [[ -f "demos/quick_demo.jl" ]]; then
        echo "Testing quick demo..."
        if timeout 30 julia --project=. demos/quick_demo.jl >/dev/null 2>&1; then
            echo -e "${GREEN}✓ Quick demo test passed${NC}"
        else
            echo -e "${YELLOW}⚠️ Quick demo test had issues${NC}"
        fi
    fi

    return 0
}

print_setup_summary() {
    local setup_type="$1"

    echo ""
    echo -e "${GREEN}================================================${NC}"
    echo -e "${GREEN}   GSICoreAnalysis.jl Setup Complete!${NC}"
    echo -e "${GREEN}================================================${NC}"
    echo ""
    echo "🎉 Installation Type: $setup_type"
    echo "📁 Installation Directory: $SCRIPT_DIR"
    echo ""
    echo "📋 Available Commands:"
    echo "  ./scripts/shell/run_gsi_analysis.sh        - Run GSI analysis"
    echo "  ./scripts/shell/run_gsi_visualization.sh   - Run with visualization"
    echo "  ./scripts/shell/run_gsi_tests.sh          - Run test suite"
    echo ""
    echo "🚀 Quick Start Examples:"
    echo "  ./scripts/shell/run_gsi_analysis.sh quick           # 10-second demo"
    echo "  ./scripts/shell/run_gsi_visualization.sh full       # Complete analysis with plots"
    echo "  ./scripts/shell/run_gsi_tests.sh minimal            # Basic functionality tests"
    echo ""
    echo "📚 Documentation:"
    echo "  • README.md - Package overview"
    echo "  • docs/ - Detailed documentation"
    echo "  • demos/ - Example scripts"
    echo ""
    echo "🔧 Environment Setup:"
    echo "  source .gsi_env                       # Load GSI environment"
    echo "  gsi-help                              # Show available commands"
    echo ""
}

# ==============================================================================
# MAIN SCRIPT
# ==============================================================================

# Default values
SETUP_TYPE="full"
AUTO_YES="false"
FORCE="false"
SKIP_TESTS="false"
SKIP_PRECOMPILE="false"
JULIA_THREADS="auto"
CHECK_ONLY="false"

# Parse arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        -h|--help)
            print_usage
            exit 0
            ;;
        -y|--yes)
            AUTO_YES="true"
            shift
            ;;
        -f|--force)
            FORCE="true"
            shift
            ;;
        --no-test)
            SKIP_TESTS="true"
            shift
            ;;
        --no-precompile)
            SKIP_PRECOMPILE="true"
            shift
            ;;
        --threads)
            JULIA_THREADS="$2"
            shift 2
            ;;
        --check-only)
            CHECK_ONLY="true"
            shift
            ;;
        full|minimal|development|visualization|hpc)
            SETUP_TYPE="$1"
            shift
            ;;
        *)
            echo -e "${RED}ERROR: Unknown option: $1${NC}" >&2
            print_usage
            exit 1
            ;;
    esac
done

# Main execution
print_banner

echo "🔧 Setup Configuration:"
echo "   • Setup Type: $SETUP_TYPE"
echo "   • Auto Accept: $AUTO_YES"
echo "   • Force Reinstall: $FORCE"
echo "   • Skip Tests: $SKIP_TESTS"
echo "   • Skip Precompile: $SKIP_PRECOMPILE"
echo "   • Julia Threads: $JULIA_THREADS"
echo ""

# System checks
if ! check_system_requirements; then
    echo -e "${RED}❌ System requirements not met${NC}"
    exit 1
fi

if ! check_julia_installation; then
    echo -e "${RED}❌ Julia installation check failed${NC}"
    exit 1
fi

if [[ "$CHECK_ONLY" == "true" ]]; then
    echo -e "${GREEN}✅ System requirements check completed successfully${NC}"
    exit 0
fi

# Confirm installation
if [[ "$AUTO_YES" != "true" ]]; then
    echo ""
    echo -e "${YELLOW}Ready to install GSICoreAnalysis.jl ($SETUP_TYPE setup)${NC}"
    read -p "Continue? (y/N): " -n 1 -r
    echo
    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        echo "Installation cancelled."
        exit 0
    fi
fi

# Run setup steps
echo ""
echo -e "${BLUE}🚀 Starting installation...${NC}"

if ! setup_julia_environment "$SETUP_TYPE"; then
    echo -e "${RED}❌ Julia environment setup failed${NC}"
    exit 1
fi

create_run_scripts_executable
setup_environment_file

if ! run_post_install_tests; then
    echo -e "${YELLOW}⚠️ Post-installation tests had issues${NC}"
    echo "The installation may still work, but consider running tests manually"
fi

print_setup_summary "$SETUP_TYPE"

echo -e "${GREEN}🎉 Setup completed successfully!${NC}"